#include <iostream>
#include <cuda_runtime.h>
#include "cublas_v2.h"
#include "../common/common.h"
#define DEVICENUM 0
#define M 2 // 2行
#define N 3 // 3列

int main(int argc, char **argv)
{
    cudaSetDevice(DEVICENUM);

    float alpha = 1.0f;
    float beta = 0.0f;

    // 声明主机端矩阵A，主机端向量X，主机端向量Y，并申请内存
    float *A = (float *)malloc(sizeof(float) * M * N);
    float *X = (float *)malloc(sizeof(float) * N);
    float *Y = (float *)malloc(sizeof(float) * M);

    initialDataIncreasing(A, M * N);
    // initialDataConstValue(X, N, 1);
    initialDataIncreasing(X, N);
    initialDataConstValue(Y, M, 1);

    printMatrix(A, M, N, "A");
    printVector(X, N, "X");
    printVector(Y, M, "Y");
    std::cout << "alpha=" << alpha << ",beta=" << beta << std::endl;

    // 声明设备端矩阵dA，设备端向量dX，设备端向量dY,并出申请内存
    float *dA, *dX, *dY;
    CHECK(cudaMalloc((void **)&dA, sizeof(float) * M * N));
    CHECK(cudaMalloc((void **)&dX, sizeof(float) * N));
    CHECK(cudaMalloc((void **)&dY, sizeof(float) * M));

    cublasHandle_t handle = 0;
    CHECK_CUBLAS(cublasCreate(&handle));

    CHECK_CUBLAS(cublasSetVector(N, sizeof(float), X, 1, dX, 1));
    CHECK_CUBLAS(cublasSetVector(M, sizeof(float), Y, 1, dY, 1));
    CHECK_CUBLAS(cublasSetMatrix(M, N, sizeof(float), A, M, dA, M));

    /* 1. 这里不使用转置，此时dA的矩阵排列如下,但显然，这并不是我们想要的结果
       1  3  5
       2  4  6 */
    // CHECK_CUBLAS(cublasSgemv(handle, CUBLAS_OP_N, M, N, &alpha, dA, M, dX, 1,
    //                          &beta, dY, 1));
    // CHECK_CUBLAS(cublasGetVector(M, sizeof(float), dY, 1, Y, 1));
    // printVector(Y, M, "no trans res");

    /* 2. 只进行转置，但不改变lda主维度
        1 2
        3 4
        5 6 */
    // CHECK_CUBLAS(cublasSgemv(handle, CUBLAS_OP_T, M, N, &alpha, dA, M, dY, 1,
    //                          &beta, dX, 1));
    // CHECK_CUBLAS(cublasGetVector(N, sizeof(float), dX, 1, X, 1));
    // printVector(X, N, "keep lda trans res");

    /* 3. 执行真正的转置
     3.1 设置dA维度为(N,M),后，其排列为：
        1 4
        2 5
        3 6
     3.2 进行转置后，dA排序为
        1 2 3
        4 5 6
    */
    CHECK_CUBLAS(cublasSgemv(handle, CUBLAS_OP_T, N, M, &alpha, dA, N, dX, 1,
                             &beta, dY, 1));
    CHECK_CUBLAS(cublasGetVector(M, sizeof(float), dY, 1, Y, 1));
    printVector(Y, M, "true trans res");

    /* 4.总结，要想去除cublas的列优先影响，需要做哪些事情？如果你暂时无法理解的话，可以先记下来
         4.1 将trans设置为CUBLAS_OP_T
         4.2 交换N,M的位置
         4.3 修改lda为N
    */

    free(A);
    free(X);
    free(Y);
    cudaFree(dA);
    cudaFree(dX);
    cudaFree(dY);

    CHECK_CUBLAS(cublasDestroy(handle));
    return 0;
}